from channels.generic.websocket import WebsocketConsumer
from channels.exceptions import StopConsumer

from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer

import json
import time
import threading
from tools.myredis import mredis


class LogoGenerationConsumer(WebsocketConsumer):
    """Logo生成结果推送的WebSocket消费者"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.user_id = None
        self.group_name = None
        self.polling_thread = None
        self.stop_polling = False
    
    def websocket_connect(self, message):
        print("Logo生成WebSocket连接建立")
        # 从URL参数或查询字符串获取用户ID
        user_id = self.scope['url_route']['kwargs'].get('user_id')
        if not user_id:
            # 如果没有用户ID，拒绝连接
            self.close()
            return
            
        self.user_id = user_id
        self.group_name = f"logo_generation_{user_id}"
        
        # 将用户加入对应的房间组
        async_to_sync(self.channel_layer.group_add)(self.group_name, self.channel_name)
        
        # 接受WebSocket连接
        self.accept()
        
        # 发送连接成功消息
        self.send(json.dumps({
            'type': 'connection_established',
            'message': 'Logo生成WebSocket连接已建立',
            'user_id': user_id
        }))
        
        # 启动Redis轮询线程
        self.start_redis_polling()

    def websocket_receive(self, message):
        """接收前端消息（可选，用于心跳检测等）"""
        try:
            data = json.loads(message['text'])
            print(f"收到Logo生成WebSocket消息: {data}")
            
            # 处理心跳检测
            if data.get('type') == 'ping':
                self.send(json.dumps({
                    'type': 'pong',
                    'timestamp': data.get('timestamp')
                }))
        except json.JSONDecodeError:
            print("无效的JSON消息")

    def websocket_disconnect(self, message):
        """WebSocket连接断开"""
        print(f"Logo生成WebSocket连接断开，用户ID: {getattr(self, 'user_id', 'unknown')}")
        
        # 停止轮询线程
        self.stop_redis_polling()
        
        if hasattr(self, 'group_name'):
            async_to_sync(self.channel_layer.group_discard)(
                self.group_name,
                self.channel_name
            )
        raise StopConsumer()

    def start_redis_polling(self):
        """启动Redis轮询线程"""
        if self.polling_thread is None or not self.polling_thread.is_alive():
            self.stop_polling = False
            self.polling_thread = threading.Thread(target=self._poll_redis_data)
            self.polling_thread.daemon = True
            self.polling_thread.start()
            print(f"启动Redis轮询线程，用户ID: {self.user_id}")

    def stop_redis_polling(self):
        """停止Redis轮询线程"""
        self.stop_polling = True
        if self.polling_thread and self.polling_thread.is_alive():
            self.polling_thread.join(timeout=2)
        print(f"停止Redis轮询线程，用户ID: {self.user_id}")

    def _poll_redis_data(self):
        """轮询Redis数据并推送给前端"""
        last_progress = None
        last_result = None
        
        while not self.stop_polling:
            try:
                # 检查进度更新
                progress_key = f"logo_progress_{self.user_id}"
                progress_data = mredis.get(progress_key)
                
                if progress_data:
                    progress_info = json.loads(progress_data)
                    current_progress = progress_info.get('progress', 0)
                    
                    # 如果进度有变化，推送给前端
                    if last_progress != current_progress:
                        try:
                            self.send(json.dumps({
                                'type': 'logo_progress',
                                'data': progress_info,
                                'timestamp': time.time()
                            }, ensure_ascii=False))
                            last_progress = current_progress
                        except Exception as send_error:
                            print(f"发送进度消息失败，用户ID: {self.user_id}, 错误: {send_error}")
                
                # 检查结果更新
                result_key = f"logo_result_{self.user_id}"
                result_data = mredis.get(result_key)
                
                if result_data:
                    result_info = json.loads(result_data)
                    result_timestamp = result_info.get('timestamp', 0)
                    
                    # 如果结果有更新，推送给前端
                    if last_result != result_timestamp:
                        try:
                            self.send(json.dumps({
                                'type': 'logo_result',
                                'data': result_info,
                                'timestamp': time.time()
                            }, ensure_ascii=False))
                            last_result = result_timestamp
                            
                            # 如果任务完成，停止轮询并断开连接
                            if result_info.get('success') is not None:
                                print(f"Logo生成任务完成，停止轮询，用户ID: {self.user_id}")
                                # 等待更长时间确保前端状态更新完成
                                time.sleep(10)
                                # 主动断开WebSocket连接
                                self.close()
                                break
                        except Exception as send_error:
                            print(f"发送结果消息失败，用户ID: {self.user_id}, 错误: {send_error}")
                
                # 每2秒轮询一次
                time.sleep(2)
                
            except Exception as e:
                print(f"Redis轮询错误，用户ID: {self.user_id}, 错误: {e}")
                time.sleep(5)  # 出错时等待更长时间

    def logo_result(self, event):
        """推送Logo生成结果给前端（保留原有方法以兼容）"""
        try:
            result_data = event.get('data', {})
            message_type = event.get('message_type', 'logo_result')
            
            # 发送结果给前端
            self.send(json.dumps({
                'type': message_type,
                'data': result_data,
                'timestamp': event.get('timestamp')
            }, ensure_ascii=False))
            
        except Exception as e:
            print(f"推送Logo结果失败: {e}")
            # 发送错误消息给前端
            self.send(json.dumps({
                'type': 'error',
                'message': '推送结果失败',
                'error': str(e)
            }))

    def logo_progress(self, event):
        """推送Logo生成进度给前端（保留原有方法以兼容）"""
        try:
            progress_data = event.get('data', {})
            
            self.send(json.dumps({
                'type': 'logo_progress',
                'data': progress_data,
                'timestamp': event.get('timestamp')
            }, ensure_ascii=False))
            
        except Exception as e:
            print(f"推送Logo进度失败: {e}")


